from contextlib import asynccontextmanager
from typing import AsyncGenerator, Callable, Any

from fastapi import Request
from sqlalchemy.exc import SQLAlchemyError, OperationalError
from sqlalchemy.ext.asyncio import AsyncSession

from typing import AsyncIterator, Optional

from sqlalchemy.ext.asyncio import AsyncEngine, async_scoped_session


class DatabaseSessionDependency:

    def __init__(self) -> None:
        self._engine: Optional[AsyncEngine] = None
        self._session: Optional[async_scoped_session] = None

    async def __call__(self) -> AsyncIterator[async_scoped_session]:
        """
        """
        async with self._session() as session:
            # await async_session.commit()
            try:
                yield session
                # 使用自动提交在某些情况下 会突发的asyncio.exceptions.CancelledError异常
                await session.commit()
            except SQLAlchemyError as ex:
                await session.rollback()
            finally:
                # 1.4.0版本的话 需要收的执行关闭，才可以 不然会一直超时
                # 1.4.19的话 这个地方可以起到关闭的作用！
                await session.close()

    async def close_on_event(self) -> None:
        """Shut down the database engine."""
        if self._engine:
            await self._engine.dispose()
            self._engine = None

    async def setup_on_event(
            self,
            _engine: AsyncEngine,
            _session: Optional[AsyncIterator[async_scoped_session]],
    ) -> None:
        '''
        该方法要使用当前类依赖项实例化的对话，在插件中完成初始化处理
        :param _engine:
        :param _session:
        :return:
        '''
        self._engine = _engine
        self._session = _session


async_session_class_v2_dependency = DatabaseSessionDependency()
